/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.lyh.api;

import org.apache.shardingsphere.driver.api.ShardingSphereDataSourceFactory;
import org.apache.shardingsphere.infra.config.algorithm.AlgorithmConfiguration;
import org.apache.shardingsphere.infra.config.props.ConfigurationPropertyKey;
import org.apache.shardingsphere.infra.config.rule.RuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.ShardingRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.rule.ShardingTableReferenceRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.rule.ShardingTableRuleConfiguration;
import org.apache.shardingsphere.sharding.api.config.strategy.sharding.StandardShardingStrategyConfiguration;

import javax.sql.DataSource;
import java.sql.*;
import java.util.*;

/**
 * 测试基于时间范围
 */
public class TestDateRangeSharding {

    public static void main(String[] args) throws SQLException {
        Collection<RuleConfiguration> ruleConfigs = new ArrayList<>();
        ruleConfigs.add(getShardingRule());
        Map<String, DataSource> dataSourceMap = new HashMap<>();
        dataSourceMap.put("ds", StaticConfiguration.createDataSource());
        DataSource dataSource = ShardingSphereDataSourceFactory.createDataSource("sharding-datasource", null,
                dataSourceMap, ruleConfigs, getProperties());
        System.out.println("--------------------------");
        createTables(dataSource);
        System.out.println("--------------------------");
        deleteData(dataSource);
        System.out.println("--------------------------");
        insertData(dataSource);
        System.out.println("--------------------------");
        queryAll(dataSource);
        System.out.println("--------------------------");
        backup(dataSource);
    }

    public static void queryAll(DataSource dataSource) throws SQLException {
        try (Connection connection = dataSource.getConnection()) {
            String sql = "SELECT * FROM t_order";
            try (PreparedStatement ps = connection.prepareStatement(sql)) {
                try (ResultSet resultSet = ps.executeQuery()) {
                    while (resultSet.next()) {
                        Long id = resultSet.getLong("id");
                        String name = resultSet.getString("name");
                        System.out.println("row id:" + id + ",name:" + name);
                        System.out.println("----^v^ 测试记录第一条，逃跑了... ^v^----");
                        break;
                    }
                }
            }
            try (PreparedStatement ps = connection.prepareStatement("SELECT count(*) as total FROM t_order")) {
                try (ResultSet resultSet = ps.executeQuery()) {
                    while (resultSet.next()) {
                        Long total = resultSet.getLong("total");
                        System.out.println("total:" + total);
                    }
                }
            }
        }
    }

    public static void backup(DataSource dataSource) throws SQLException {
        try (Connection connection = dataSource.getConnection()) {
            String sql = "INSERT INTO t_order_back SELECT * FROM t_order";
            try (PreparedStatement ps = connection.prepareStatement(sql)) {
                System.out.println("批量备份：" + ps.execute());
            }
        }
    }

    public static ShardingRuleConfiguration getShardingRule() {
        String USER_ALGORITHM = "CUSTOM_DATE_RANGE";
        ShardingRuleConfiguration shardingRuleConfiguration = new ShardingRuleConfiguration();
        int size = TimeRouter.of("ds.t_order").getTables().size();
        {


            String actualDataNodes = TimeRouter.of("ds.t_order").getTablesAsString();
            ShardingTableRuleConfiguration ruleConfiguration2022 = new ShardingTableRuleConfiguration("t_order", actualDataNodes);
            ruleConfiguration2022.setTableShardingStrategy(new StandardShardingStrategyConfiguration("create_time", USER_ALGORITHM));

            shardingRuleConfiguration.getTables().add(ruleConfiguration2022);
            Properties userShardingAlgorithmProperties = new Properties();
            userShardingAlgorithmProperties.put("sharding-count", size);
            AlgorithmConfiguration algorithmConfiguration = new AlgorithmConfiguration("CUSTOM_DATE_RANGE",
                    userShardingAlgorithmProperties);
            shardingRuleConfiguration.getShardingAlgorithms().put(USER_ALGORITHM, algorithmConfiguration);
        }
        {
            String actualDataNodes = TimeRouter.of("ds.t_order_back").getTablesAsString();
            ShardingTableRuleConfiguration ruleConfiguration2022 = new ShardingTableRuleConfiguration("t_order_back", actualDataNodes);
            ruleConfiguration2022.setTableShardingStrategy(new StandardShardingStrategyConfiguration("create_time", USER_ALGORITHM));
            shardingRuleConfiguration.getTables().add(ruleConfiguration2022);
            Properties userShardingAlgorithmProperties = new Properties();
            userShardingAlgorithmProperties.put("sharding-count", size);
            AlgorithmConfiguration algorithmConfiguration = new AlgorithmConfiguration("CUSTOM_DATE_RANGE",
                    userShardingAlgorithmProperties);
            shardingRuleConfiguration.getShardingAlgorithms().put(USER_ALGORITHM, algorithmConfiguration);
        }
        shardingRuleConfiguration.getBindingTableGroups().add(
                new ShardingTableReferenceRuleConfiguration("foo", "t_order,t_order_back"));
        return shardingRuleConfiguration;
    }

    private static Properties getProperties() {
        Properties props = new Properties();
        props.put(ConfigurationPropertyKey.SQL_SHOW.getKey(), "true");
        return props;
    }

    private static void deleteData(DataSource dataSource) throws SQLException {
        String sql = "DELETE FROM t_order";
        try (Connection connection = dataSource.getConnection();
             PreparedStatement ps = connection.prepareStatement(sql)) {
            System.out.println(ps.execute());
        }
    }

    private static void insertData(DataSource dataSource) throws SQLException {
        String sql = "insert into t_order(id,name,create_time) values (?,?,?)";
        boolean batchInsert = true;
        if (batchInsert) {
            try (Connection connection = dataSource.getConnection();
                 PreparedStatement ps = connection.prepareStatement(sql)) {
                for (int i = 1; i <= 10; i++) {
                    ps.setLong(1, i);
                    ps.setString(2, "测试消息-" + i);
                    ps.setDate(3, MockData.randomDate());
                    ps.addBatch();
                }
                int[] results = ps.executeBatch();
                System.out.println("批量插入:" + Arrays.toString(results));
            }
        } else {
            for (int i = 1; i <= 10; i++) {
                try (Connection connection = dataSource.getConnection();
                     PreparedStatement ps = connection.prepareStatement(sql)) {
                    ps.setLong(1, i);
                    ps.setString(2, "测试消息-" + i);
                    ps.setDate(3, MockData.randomDate());
                    System.out.println(ps.execute());
                }
            }
        }
    }

    private static void createTables(DataSource dataSource) throws SQLException {
        try (Connection connection = dataSource.getConnection(); Statement statement = connection.createStatement()) {
            String sql = "CREATE TABLE IF NOT EXISTS t_order(id BIGINT, name VARCHAR(100) NOT NULL, create_time TIMESTAMP NOT NULL, PRIMARY KEY (id))";
            System.out.println("创建表:" + statement.execute(sql));
            String sql_0 = "CREATE TABLE IF NOT EXISTS t_order_back(id BIGINT, name VARCHAR(100) NOT NULL, create_time TIMESTAMP NOT NULL, PRIMARY KEY (id))";
            System.out.println("创建表:" + statement.execute(sql_0));
        }
    }
}